#!/usr/bin/env python2.7
#-*- coding=utf-8 -*-

from ..cstructures import BigendCStruct
from .base import NetworkData
from .LinkLayer import _formatmac_
import re
import struct

def _formatip_(ipdata, version = None, flags = None):
    if version is None:
        version = 4
    if version == 4:
        ipsec = []
        for i in (24, 16, 8, 0):
            ipsec.append(str((ipdata >> i) & 0xff))
        return '.'.join(ipsec)
    elif version == 6:
        if flags is None:
            flags = 0
        ipsec = []
        for i in ipdata:
            section = '%04x' % i
            if flags & 1:
                section = re.sub(r'^0+', '', section)
                if section == '':
                    section = '0'
            ipsec.append(section)
        ip = ':'.join(ipsec)
        if flags & 2:
            ip = re.sub(r'(^|:)[0:]+(:|$)', '::', ip, 1)
        return ip
    return None

class IPv4(NetworkData):
    def __init__(self, data):
        super(IPv4, self).__init__(data)
        #split version and hdrlen
        self.version = self.verlen >> 0x4
        self.hdrlen = (self.verlen & 0xf) * 0x4
        optlen = self.hdrlen - 0x14
        if optlen > 0:
            self.options = self.restdata[:optlen]
            if len(self.options) < optlen:
                raise BufferError('Data buffer is not long enough!')
            self.restdata = self.restdata[optlen:]
        #split DSCP and ECN
        self.dscp = self.tos >> 0x2
        self.ecn = self.tos & 0x2
        #split Flags and Fragment Offset
        self.flags = self.flagoffset >> 0xd
        self.offset = self.flagoffset & 0xd
        #get the rest of data
        datalen = self.ttl - 0x14 - optlen
        self.restdata = self.restdata[:datalen]
        self.nextlayer = self.protocol

    def getnames(self):
        names = super(IPv4, self).getnames()
        #split version and hdrlen
        names[0] = 'version'
        names.insert(1, 'hdrlen')
        #split DSCP and ECN
        names[2] = 'dscp'
        names.insert(3, 'ecn')
        #split Flags and Fragment Offset
        names[6] = 'flags'
        names.insert(7, 'offset')

    def verify(self):
        checksum = (self.verlen << 8) + self.tos
        checksum += self.tlen
        checksum += self.indentification
        checksum += self.flagoffset
        checksum += (self.ttl << 8) + self.protocol
        checksum += self.checksum
        checksum += (self.srcip >> 16) + (self.srcip & 0xffff)
        checksum += (self.dstip >> 16) + (self.dstip & 0xffff)
        checksum = (checksum >> 16) + (checksum & 0xffff)
        if (checksum & 0xffff) == 0xffff:
            return True
        return False

    def formatip(self):
        srcip = _formatip_(self.srcip)
        dstip = _formatip_(self.dstip)
        return srcip, dstip

    _fields_ = [
            ('verlen', 'uByte'),
            ('tos', 'uByte'),
            ('tlen', 'WORD'),
            ('indentification', 'WORD'),
            ('flagoffset', 'WORD'),
            ('ttl', 'uByte'),
            ('protocol', 'uByte'),
            ('checksum', 'WORD'),
            ('srcip', 'DWORD'),
            ('dstip', 'DWORD')
        ]

class IPv6(NetworkData):
    def __init__(self, data):
        super(IPv6, self).__init__(data)
        self.version = self.vertrafflow >> 28
        self.traffic = (self.vertrafflow >> 20) & 0xff
        self.flowlabel = self.vertrafflow & 0xfffff
        self.nextlayer = self.nextheader

    def getnames(self):
        names = super(IPv6, self).getnames()
        names.remove('vertrafflow')
        newnames = ['version', 'traffic', 'flowlabel']
        newnames.extend(names)
        return newnames

    def formatip(self):
        srcip = _formatip_(self.srcip, 6, 3)
        dstip = _formatip_(self.dstip, 6, 3)
        return srcip, dstip

    _fields_ = [
            ('vertrafflow', 'DWORD'),
            ('payloadlen', 'WORD'),
            ('nextheader', 'uByte'),
            ('hoplimit', 'uByte'),
            ('srcip', 'WORD', 8),
            ('dstip', 'WORD', 8)
        ]

class ARP(NetworkData):
    def __init__(self, data):
        super(ARP, self).__init__(data)
        hds = self.hardsize
        pts = self.protosize
        addrlen = (hds + pts) * 2
        addrdata = data[self.datalen:self.datalen + addrlen]
        if len(addrdata) < addrlen:
            raise BufferError('Data buffer is not long enough!')
        self.restdata = data[self.datalen + addrlen:]
        self.datalen += addrlen
        self.sndmac = struct.unpack('>%dB' % hds, addrdata[:hds])
        self.tarmac = struct.unpack('>%dB' % hds, addrdata[hds + pts:hds * 2 + pts])
        if pts == 4:
            self.sndpa = struct.unpack('>I', addrdata[hds:hds + pts])[0]
            self.tarpa = struct.unpack('>I', addrdata[hds * 2 + pts:])[0]
            self.ptv = 4
        elif pts == 16:
            self.sndpa = struct.unpack('>%dH' % pts/2, addrdata[hds:hds + pts])
            self.tarpa = struct.unpack('>%dH' % pts/2, addrdata[hds * 2 + pts:])
            self.ptv = 6
        else:
            raise TypeError('Protocol Size is %d. Connot recognize!' % pts)

    def getnames(self):
        names = super(ARP, self).getnames()
        names.extend('sndmac', 'sndpa', 'tarmac', 'tarpa')
        return names

    def formataddr(self):
        sndmac = _formatmac_(self.sndmac)
        tarmac = _formatmac_(self.tarmac)
        sndpa = _formatip_(self.sndpa, self.ptv, 3)
        tarpa = _formatip_(self.tarpa, self.ptv, 3)
        return sndmac, sndpa, tarmac, tarpa

    _fields_ = [
        ('hardtype', 'WORD'),
        ('prototype', 'WORD'),
        ('hardsize', 'uByte'),
        ('protosize', 'uByte'),
        ('opcode', 'WORD'),
    ]

RARP = ARP

layerdict = {
        0x800:IPv4,
        0x806:ARP,
        0x8035:RARP,
        0x86dd:IPv6
    }
